#! /usr/bin/env python
#coding=utf-8
from ctypes import *
from datetime import datetime
from matplotlib import pyplot as plt
import numpy as np

usb_id = 0
dev_addr = 0x80
resistor=0.1 #ohm
battery_capacity=3.4*3*1000 #mAh

ch341 = windll.LoadLibrary("./CH341DLLA64.DLL")
if ch341.CH341OpenDevice(usb_id) == -1:
    print("USB CH341 Open Failed!")
    exit(1)
ch341.CH341SetStream(usb_id, 0x81)

def ina219config(dat):
    """change config register in ina219

    Args:
        dat (byte[2]): new conf data
    """
    addr=0x00
    obuf = (c_byte * 4)()
    ibuf = (c_byte * 1)()
    obuf[0] = dev_addr
    obuf[1] = addr
    obuf[2] = dat[0] & 0xff
    obuf[3] = dat[1] & 0xff
    ch341.CH341StreamI2C(usb_id, 4, obuf, 0, ibuf)

def sample_data(sample_time=10):
    """sample current data in given time length

    Args:
        sample_time (int, optional): time you want to sample. Defaults to 10.

    Returns:
        data_t, data_c: time & current data lists
    """
    data_addr=0x01
    samplerate_hz=1770
    rec  = (c_ubyte * 2)()
    ibuf = (c_ubyte * 2)()
    rec[0] = dev_addr
    rec[1] = data_addr & 0xff
    data_t=[]
    data_c=[]
    current_old=0
    # print("calc time:",round(sample_num/samplerate_hz,2),"s")
    # print("sample num:",sample_num)
    starttime=datetime.timestamp(datetime.now())
    while True:
        #read data
        ch341.CH341StreamI2C(usb_id, 2, rec, 2, ibuf)
        # current=round((ibuf[0] << 8 | ibuf[1])*0.01,2)
        current=round((ibuf[0] << 8 | ibuf[1])/resistor*0.01,2)
        if(current>3.2*1024):
            current=current-6.4*1024
        # print(current,"mV")
        if current == current_old:
            continue
        current_old=current
        timenow=datetime.timestamp(datetime.now())-starttime
        if(timenow>sample_time):
            break
        #update 2 arrays
        data_t.append(timenow)
        data_c.append(current)
        if(len(data_t)%1000==0):
            print("sample num:",len(data_t))
    ch341.CH341CloseDevice(usb_id)
    endtime=datetime.timestamp(datetime.now())
    print("sample time:",round(endtime-starttime,2),"s")
    print("sample num:",len(data_t))
    true_sample_rate=round(len(data_t)/(endtime-starttime),2)
    print("sample rate:",true_sample_rate,"Hz")
    #save data to csv
    datanp=np.array([data_t,data_c])
    datanp=datanp.T
    fname="./data/"
    fname+=datetime.now().strftime("%y%m%d_%H%M%S")+".csv"
    np.savetxt(fname,datanp,delimiter=",",fmt="%f")
    print("data saved to",fname)
    return data_t,data_c

def plotdata(data_t,data_c):
    plt.title("current vs. time")
    plt.xlabel("time(s)")
    plt.ylabel("current(mA)")
    plt.ylim(0,max(data_c)*1.1)
    plt.plot(data_t,data_c)
    plt.show()

#calc energy, return mA*s
def calc_energy_consumption(data_t,data_c):
    energy=0
    for i in range(len(data_t)-1):
        energy+=((data_c[i]+data_c[i+1])/2)*(data_t[i+1]-data_t[i])
    return energy

def calc_average_current(data_c):
    sumcurrent=sum(data_c)
    return round(sumcurrent/len(data_c),2)

def read_2byte(addr):
    obuf = (c_ubyte * 2)()
    ibuf = (c_ubyte * 2)()
    obuf[0] = dev_addr
    obuf[1] = addr
    ch341.CH341StreamI2C(usb_id, 2, obuf, 2, ibuf)
    return ibuf[0] & 0xff,ibuf[1] & 0xff


def read_current():
    data_addr=0x01
    rec  = (c_ubyte * 2)()
    ibuf = (c_ubyte * 2)()
    rec[0] = dev_addr
    rec[1] = data_addr & 0xff
    ch341.CH341StreamI2C(usb_id, 2, rec, 2, ibuf)
    ch341.CH341CloseDevice(usb_id)
    current=round((ibuf[0] << 8 | ibuf[1])/resistor*0.01,2)
    if(current>3.2*1024):
        current=current-6.4*1024
    return current #mA


if __name__ == "__main__":
    # print(read_current())
    # exit()

    ina219config([0x21,0x8d])
    rec=read_2byte(0x00)
    print(hex(rec[0] << 8 | rec[1]))

    data_t,data_c=sample_data(5)
    plotdata(data_t,data_c)

    averagec=calc_average_current(data_c)
    print("average current:",averagec,"mA")
    print("max current:",max(data_c),"mA")

    energy=calc_energy_consumption(data_t,data_c) #mAs
    print("energy consumption:",round(energy/3600,4),"mAh")
    print("battery capacity:",battery_capacity,"mAh")

    timelen=data_t[-1]-data_t[0]
    estimated_battery_life=battery_capacity*3600/energy*timelen #s
    print("estimated battery life:",round(estimated_battery_life/3600,2),"h ==",round(estimated_battery_life/3600/24,2),"d")